package cn.monkey.state.disruptor;

import cn.monkey.state.core.Queue;
import cn.monkey.state.scheduler.disruptor.EventWrapper;
import com.lmax.disruptor.*;
import com.lmax.disruptor.dsl.ProducerType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;

public class RingBufferQueue<E> implements Queue<E> {

    private static final Logger log = LoggerFactory.getLogger(RingBufferQueue.class);

    private final RingBuffer<EventWrapper<E>> ringBuffer;

    private final Sequence sequence;

    private final SequenceBarrier sequenceBarrier;

    public RingBufferQueue(ProducerType producerType,
                           int bufferSize,
                           WaitStrategy waitStrategy) {
        this.ringBuffer = RingBuffer.create(producerType, EventWrapper::new, bufferSize, waitStrategy);
        this.sequenceBarrier = this.ringBuffer.newBarrier();
        this.sequence = new Sequence(Sequencer.INITIAL_CURSOR_VALUE);
    }


    @Override
    public boolean offer(E e) {
        long next;
        try {
            next = this.ringBuffer.tryNext();
        } catch (InsufficientCapacityException ex) {
            return false;
        }
        try {
            EventWrapper<E> wrapper = this.ringBuffer.get(next);
            wrapper.event = e;
            return true;
        } finally {
            this.ringBuffer.publish(next);
        }

    }

    @Override
    public E poll() {
        long nextSequence = this.sequence.get() + 1L;
        try {
            long availableSequence = this.sequenceBarrier.waitFor(nextSequence);
            if (nextSequence <= availableSequence) {
                EventWrapper<E> eventWrapper = this.ringBuffer.get(nextSequence);
                this.sequence.incrementAndGet();
                return eventWrapper.event;
            }
        } catch (AlertException | InterruptedException | TimeoutException e) {
            log.error("message poll error:\n", e);
            return null;
        }
        return null;
    }

    @Override
    public List<E> poll(int n) {
        List<E> events = new ArrayList<>(n);
        long ringBufferSize = this.ringBuffer.getBufferSize();
        long nextSequence = this.sequence.get() + 1L;

        try {
            long availableSequence = this.sequenceBarrier.waitFor(nextSequence);
            long endOfBatchSequence = Math.min(nextSequence + ringBufferSize, availableSequence);
            do {
                EventWrapper<E> eventWrapper = this.ringBuffer.get(nextSequence);
                events.add(eventWrapper.event);
                nextSequence++;
            } while (nextSequence <= availableSequence);
            this.sequence.set(endOfBatchSequence);
            return events;
        } catch (AlertException | InterruptedException | TimeoutException e) {
            log.error("message poll error:\n", e);
            return events;
        }
    }
}
